package rename_table

import (
	"bytes"
	"fmt"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/format"
	"strings"
)

// SQLRestorer 用于将AST转换回SQL字符串
type SQLRestorer struct {
	buffer *bytes.Buffer
}

// NewSQLRestorer 创建新的SQL恢复器
func NewSQLRestorer() *SQLRestorer {
	return &SQLRestorer{
		buffer: &bytes.Buffer{},
	}
}

// RestoreSQL 将AST节点转换为SQL字符串
func (r *SQLRestorer) RestoreSQL(node ast.Node) (string, error) {
	r.buffer.Reset()

	ctx := format.NewRestoreCtx(format.DefaultRestoreFlags, r.buffer)
	err := node.Restore(ctx)
	if err != nil {
		return "", fmt.Errorf("恢复SQL失败: %v", err)
	}

	return r.buffer.String(), nil
}

// RenameTablesInSQLWithRestore 解析SQL语句，重命名表名，并返回修改后的SQL
func RenameTablesInSQLWithRestore(sql string, suffix string) (string, error) {
	p := parser.New()

	// 解析SQL语句
	stmtNodes, _, err := p.Parse(sql, "", "")
	if err != nil {
		return "", fmt.Errorf("解析SQL失败: %v", err)
	}

	// 创建访问者
	visitor := NewTableRenameVisitor(suffix)

	// 遍历所有语句节点并修改
	for _, stmt := range stmtNodes {
		stmt.Accept(visitor)
	}

	// 将修改后的AST转换回SQL字符串
	restorer := NewSQLRestorer()
	var resultSQL string

	if len(stmtNodes) == 1 {
		// 单个语句
		resultSQL, err = restorer.RestoreSQL(stmtNodes[0])
		if err != nil {
			return "", err
		}
	} else {
		// 多个语句，用分号连接
		var sqls []string
		for _, stmt := range stmtNodes {
			sqlStr, err := restorer.RestoreSQL(stmt)
			if err != nil {
				return "", err
			}
			sqls = append(sqls, sqlStr)
		}
		resultSQL = strings.Join(sqls, "; ")
	}

	return resultSQL, nil
}
